# encoding=utf-8
import argparse
import hashlib
import os
import tempfile
import time
import logging
import traceback
import cv2
import fastapi
import uvicorn
from asyncio import Lock
from fastapi import FastAPI
from starlette.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from NNUnetV1Helper import NNUnetV1MultiHelper
from NNUnetTools import AppConfig, set_logger

parser = argparse.ArgumentParser()
parser.add_argument("--path", type=str)
args = parser.parse_args()
if args.path is None:
    logging.error("run path is empty")
    exit(-1)

run_path = args.path

os.makedirs(AppConfig.CACHE_PATH, exist_ok=True)
os.makedirs(AppConfig.LOGGER_PATH, exist_ok=True)
set_logger()

optovue_five_model = NNUnetV1MultiHelper(run_path)

def generate_time_mills():
    return int(round(time.time() * 1000))


def generate_task_id(name: str):
    str_data = name + str(generate_time_mills())
    encoded = hashlib.sha256(str_data.encode()).hexdigest()
    return encoded


app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"]
)

async_lock = Lock()


class OptovueFiveController:
    @staticmethod
    @app.post("/run/optovue/five/opti")
    async def run(file: fastapi.UploadFile):
        # save to tmp, use PIL to decode
        temp_dir_name = generate_task_id("hello")
        abs_workdir = os.path.join(AppConfig.CACHE_PATH, temp_dir_name)
        abs_user_path = os.path.join(abs_workdir, "workdir", "user_input")
        os.makedirs(abs_user_path, exist_ok=True)

        logging.info("save user path: {}".format(abs_user_path))

        temp_file = os.path.basename(tempfile.NamedTemporaryFile(suffix=".png").name)
        abs_file_path = os.path.join(abs_user_path, temp_file)
        with open(abs_file_path, 'wb') as abs_file:
            abs_file.write(await file.read())

        logging.info("write done file: {}".format(abs_file_path))

        # run infer
        async with async_lock:
            try:
                start_time = time.perf_counter()
                out = optovue_five_model.infer(abs_file_path)
                out_path = os.path.join(abs_workdir, "workdir", "out_raw_data")
                os.makedirs(out_path, exist_ok=True)
                cv2.imwrite(os.path.join(out_path, temp_file), out)
                end_time = time.perf_counter()
                logging.info("infer cost: {}s".format(end_time - start_time))
            except Exception as e:
                logging.exception(traceback.format_exc())
                return JSONResponse(
                    status_code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR,
                    content={
                        'msg': traceback.format_exc()
                    }
                )

        # return
        dst_name = os.path.join(abs_workdir, "workdir", "out_raw_data", temp_file)
        return fastapi.responses.FileResponse(dst_name)


if __name__ == '__main__':
    uvicorn.run(app=app, host="0.0.0.0", port=30002)
